"""
CogView client for image generation using ZhipuAI API.
"""

import os
import aiohttp
import base64
from typing import Optional
from zhipuai import ZhipuAI
import logging

logger = logging.getLogger(__name__)

class CogViewClient:
    """Client for interacting with CogView image generation API."""
    
    def __init__(self, api_key: Optional[str] = None):
        """Initialize the CogView client.
        
        Args:
            api_key: ZhipuAI API key. If not provided, will try to get from environment.
        """
        self.api_key = api_key or os.getenv("ZHIPU_API_KEY")
        if not self.api_key:
            raise ValueError("ZHIPU_API_KEY must be provided either as parameter or environment variable")
        
        self.client = ZhipuAI(api_key=self.api_key)
    
    async def generate_image(self, prompt: str, model: str = "cogview-3-flash", size: str = "1024x1024") -> str:
        """Generate an image using CogView.
        
        Args:
            prompt: Text prompt describing the image to generate
            model: CogView model to use
            size: Image size in format "widthxheight" (e.g., "1024x1024")
            
        Returns:
            URL of the generated image
        """
        try:
            logger.info(f"Generating image with model {model}, size {size} and prompt: {prompt}")
            
            response = self.client.images.generations(
                model=model,
                prompt=prompt,
                size=size
            )
            
            if not response.data or len(response.data) == 0:
                raise Exception("No image data returned from CogView API")
            
            image_url = response.data[0].url
            logger.info(f"Image generated successfully: {image_url}")
            
            return image_url
            
        except Exception as e:
            logger.error(f"Error generating image: {str(e)}")
            raise Exception(f"Failed to generate image: {str(e)}")
    
    async def download_image(self, image_url: str) -> str:
        """Download image from URL and return as base64 encoded string.
        
        Args:
            image_url: URL of the image to download
            
        Returns:
            Base64 encoded image data
        """
        try:
            logger.info(f"Downloading image from: {image_url}")
            
            async with aiohttp.ClientSession() as session:
                async with session.get(image_url) as response:
                    if response.status != 200:
                        raise Exception(f"Failed to download image: HTTP {response.status}")
                    
                    image_bytes = await response.read()
                    image_b64 = base64.b64encode(image_bytes).decode('utf-8')
                    
                    logger.info("Image downloaded and encoded successfully")
                    return image_b64
                    
        except Exception as e:
            logger.error(f"Error downloading image: {str(e)}")
            raise Exception(f"Failed to download image: {str(e)}")
